package middle;

import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;
import util.PrintUtils;

public class Solution_347 {

    public static void main(String[] args) {
        Solution_347 model = new Solution_347();
        PrintUtils.arrayInt(model.topKFrequent(new int[]{1,1,1,2,2,3}, 2));
//        PrintUtils.arrayInt(model.topKFrequent(new int[]{1}, 1));
    }
    public int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> map = new HashMap<>();
        for (int num : nums) {
            map.put(num, map.getOrDefault(num, 0) + 1);
        }
        int index = 0;
        int[][] counts = new int[map.keySet().size()][2];
        for (Integer num : map.keySet()) {
            counts[index++] = new int[]{num, map.get(num)};
        }

        index = -1;
        int begin = 0, end = counts.length - 1;
        do {
            if (index < k) {
                begin = index + 1;
            } else {
                end = index - 1;
            }
            index = find(counts, begin, end);
        } while (index == k);
        int[] array = new int[index + 1];
        for (int i = 0; i <= index; i++) {
            array[i] = counts[index][0];
        }
        return array;
    }

    private int find(int[][] nums, int i, int j) {
        if (i == j) {
            return i;
        }
        int begin = i + 1, end = j;
        while (begin < end) {
            while (begin < end && nums[end][1] > nums[i][1]) {
                end--;
            }
            while (begin < end && nums[begin][1] <= nums[i][1]) {
                begin++;
            }
            if (begin < end) {
                swap(nums, begin++, end--);
            }
        }
        swap(nums, i, begin);
        return begin;
    }

    private void swap(int[][] nums, int i, int j) {
        int[] temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }

    public int[] topKFrequent_1(int[] nums, int k) {
        Map<Integer, Integer> map = new HashMap<>();
        for (int num : nums) {
            map.put(num, map.getOrDefault(num, 0) + 1);
        }
        PriorityQueue<int[]> queue =  new PriorityQueue<>(Comparator.comparingInt(a -> a[1]));
        for (Integer num : map.keySet()) {
            queue.add(new int[]{num, map.get(num)});
            if (queue.size() > k) {
                queue.poll();
            }
        }
        int index = 0;
        int[] array = new int[queue.size()];
        while (!queue.isEmpty()) {
            array[index++] = queue.poll()[0];
        }
        return array;
    }
}
